{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import os\n",
    "import sys\n",
    "import panel as pn\n",
    "import numpy as np\n",
    "import pyvista as pv\n",
    "pv.set_plot_theme(\"document\")\n",
    "import glob\n",
    "from matplotlib.colors import ListedColormap\n",
    "from omegaconf import OmegaConf\n",
    "from torch_geometric.data import Data\n",
    "import random\n",
    "import torch\n",
    "\n",
    "pn.extension('vtk')\n",
    "os.system('/usr/bin/Xvfb :99 -screen 0 1024x768x24 &')\n",
    "os.environ['DISPLAY'] = ':99'\n",
    "os.environ['PYVISTA_OFF_SCREEN'] = 'True'\n",
    "os.environ['PYVISTA_USE_PANEL'] = 'True'\n",
    "\n",
    "DIR = os.path.dirname(os.getcwd())\n",
    "sys.path.append(DIR)\n",
    "\n",
    "from torch_points3d.datasets.panoptic.scannet import ScannetDataset, ScannetPanoptic\n",
    "from torch_points3d.datasets.segmentation.scannet import Scannet, SCANNET_COLOR_MAP\n",
    "from torch_points3d.datasets.segmentation import IGNORE_LABEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "VIZ_REL_DIR = \"outputs/2020-07-13/17-57-38/viz\"\n",
    "VIZ_DIR = os.path.join(DIR, VIZ_REL_DIR)\n",
    "all_viz_data = glob.glob(os.path.join(VIZ_DIR,'*.pt'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sample(data, sample_idx):\n",
    "    sample_mask = data.batch == sample_idx\n",
    "    out_data = Data(pos=data.pos[sample_mask], instance_labels=data.instance_labels[sample_mask], semantic_pred=data.semantic_pred[sample_mask], y = data.y[sample_mask], vote=data.vote[sample_mask])\n",
    "    pos_instances = []\n",
    "    vote_instances = []\n",
    "    for i,cl in enumerate(data.clusters):\n",
    "        instance_sample = data.batch[cl[0]]\n",
    "        if instance_sample != sample_idx:\n",
    "            continue\n",
    "        if data.cluster_type[i] == 0:\n",
    "            pos_instances.append(data.pos[cl])\n",
    "        else:\n",
    "            vote_instances.append(data.pos[cl])\n",
    "    out_data.pos_instances = pos_instances\n",
    "    out_data.vote_instances =vote_instances\n",
    "    return out_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def buil_cmap():\n",
    "    mapping = np.linspace(IGNORE_LABEL, len(Scannet.VALID_CLASS_IDS)+1, 256)\n",
    "    newcolors = np.zeros((256, 3))\n",
    "    for i, raw_label in enumerate(Scannet.VALID_CLASS_IDS):\n",
    "        newcolors[mapping >= i-0.5] = np.asarray(Scannet.SCANNET_COLOR_MAP[raw_label]) / 255.\n",
    "    return ListedColormap(newcolors)\n",
    "cmap = buil_cmap()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_idx = -1\n",
    "sample_idx = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.load(all_viz_data[file_idx])\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.load(all_viz_data[file_idx])\n",
    "def load_random_sample(event):\n",
    "    i = np.random.randint(0, data.batch.max()+1)\n",
    "    sample = get_sample(data, i)\n",
    "    \n",
    "    pl1 = pv.Plotter(notebook=True)\n",
    "    pl2 = pv.Plotter(notebook=True)\n",
    "    pl3 = pv.Plotter(notebook=True)\n",
    "    pl4 = pv.Plotter(notebook=True)\n",
    "    pl5 = pv.Plotter(notebook=True)\n",
    "    \n",
    "    # Color by points with a label\n",
    "    mask = sample.instance_labels == 0\n",
    "    pl1.add_points(sample.pos[mask].numpy(), color=\"gray\", opacity=0.4) \n",
    "    pl2.add_points(sample.pos[mask].numpy(), color=\"gray\", opacity=0.4) \n",
    "    pl4.add_points(sample.pos[mask].numpy(), color=\"gray\", opacity=0.4) \n",
    "    pl5.add_points(sample.pos[mask].numpy(), color=\"gray\", opacity=0.4) \n",
    "    \n",
    "    num_instances = sample.instance_labels.max()\n",
    "    for i in range(1,num_instances+1):\n",
    "        instance_mask = sample.instance_labels == i\n",
    "        point_cloud = pv.PolyData(sample.pos[instance_mask].numpy())\n",
    "        labels = sample.y[instance_mask].numpy()\n",
    "        point_cloud['label'] = labels\n",
    "        pl1.add_points(point_cloud,cmap=cmap, clim=[-1, len(Scannet.VALID_CLASS_IDS)+1])\n",
    "    \n",
    "    proposed_instances = sample.pos_instances\n",
    "    for instance in proposed_instances:\n",
    "        color = [random.random(),random.random(),random.random()]\n",
    "        pl2.add_points(instance.numpy(), color=color,) \n",
    "    \n",
    "    proposed_instances = sample.vote_instances\n",
    "    for instance in proposed_instances:\n",
    "        color = [random.random(),random.random(),random.random()]\n",
    "        pl4.add_points(instance.numpy(), color=color,) \n",
    "    \n",
    "    point_cloud = pv.PolyData(sample.pos.numpy())\n",
    "    point_cloud['label'] = sample.semantic_pred.numpy()\n",
    "    pl3.add_points(point_cloud,cmap=cmap, clim=[-1, len(Scannet.VALID_CLASS_IDS)+1]) \n",
    "    \n",
    "    instance_mask = sample.instance_labels != 0\n",
    "    vote_centre = sample.pos[instance_mask] + sample.vote[instance_mask]\n",
    "    pl5.add_points(vote_centre.numpy(), color='red')\n",
    "\n",
    "    pan1.object = pl1.ren_win\n",
    "    pan2.object = pl2.ren_win\n",
    "    pan3.object = pl3.ren_win\n",
    "    pan4.object = pl4.ren_win\n",
    "    pan5.object = pl5.ren_win"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl1 = pv.Plotter(notebook=True)\n",
    "pl2 = pv.Plotter(notebook=True)\n",
    "pl3 = pv.Plotter(notebook=True)\n",
    "pl4 = pv.Plotter(notebook=True)\n",
    "pl5 = pv.Plotter(notebook=True)\n",
    "pan1 = pn.panel(pl1.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=True,)\n",
    "pan2 = pn.panel(pl2.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=True,)\n",
    "pan3 = pn.panel(pl3.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=True,)\n",
    "pan4 = pn.panel(pl4.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=True,)\n",
    "pan5 = pn.panel(pl5.ren_win, sizing_mode='scale_both', aspect_ratio=1,orientation_widget=True,)\n",
    "button = pn.widgets.Button(name='Load new model', button_type='primary')\n",
    "button.on_click(load_random_sample)\n",
    "pn.Row(\n",
    "    pn.Column('## Scannet vizualise',button),\n",
    "    pn.Column(pn.Column(pan1,'Ground truth'),pn.Column(pan2, 'Pos based predictions')),\n",
    "    pn.Column(pn.Column(pan3,'Semantic predictions'), pn.Column(pan4,'Vote clusters')),\n",
    "    pn.Column(pn.Column(pan5,'Votes'))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploration of the clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_points_kernels import region_grow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.load(all_viz_data[file_idx])\n",
    "i = np.random.randint(0, data.batch.max()+1)\n",
    "sample = get_sample(data, i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.unique(sample.instance_labels)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clusters = region_grow(sample.pos.cuda() + sample.vote.cuda(),sample.semantic_pred.cuda(),torch.zeros(sample.pos.shape[0]).long().cuda(),ignore_labels=[-1,0,1], radius=0.075,min_cluster_size=32,nsample=16)\n",
    "data.clusters = [cl.cpu() for cl in clusters]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(clusters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_points = 0\n",
    "for cl in data.clusters:\n",
    "    num_points += len(cl)\n",
    "num_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = sample.instance_labels == 0\n",
    "pl = pv.Plotter(notebook=True)\n",
    "pl.add_points(sample.pos[mask].numpy(), color=\"gray\")\n",
    "for cl in clusters:\n",
    "    color = [random.random(),random.random(),random.random()]\n",
    "    pl.add_points(sample.pos[cl].numpy(), color=color,) \n",
    "pn.panel(pl.ren_win,  aspect_ratio=1,orientation_widget=True,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clusters = region_grow(data.pos.cpu(),data.semantic_pred.cpu(),data.batch.cpu(),ignore_labels=[-1,0,1],\n",
    "                radius=0.075,min_cluster_size=32,nsample=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = data.instance_labels == 0\n",
    "pl = pv.Plotter(notebook=True)\n",
    "pl.add_points(data.pos[mask].numpy(), color=\"gray\")\n",
    "for cl in clusters:\n",
    "    color = [random.random(),random.random(),random.random()]\n",
    "    pl.add_points(data.pos[cl].numpy(), color=color,) \n",
    "pn.panel(pl.ren_win,  aspect_ratio=1,orientation_widget=True,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.unique(data.semantic_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
